from data.load_data import CustomDataset
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
from model.ddpm import train,test, generate
import torch
PICTURE_SIZE=128 # 这个要和其他地方统一

transform = transforms.Compose([
    transforms.Resize((PICTURE_SIZE, PICTURE_SIZE)),  # 调整图片大小
    transforms.ToTensor(),           # 将图片转换为Tensor
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 归一化
])

root_dir = r'data/butterfly'  # 图片文件夹路径
dataset = CustomDataset(root_dir, transform=transform)
dataloader = DataLoader(dataset, batch_size=128, shuffle=True)


# TODO:应该划分数据集
train(data_loader=dataloader, epoch=10, device='cuda' if torch.cuda.is_available() else 'cpu')

test(data_loader=dataloader, device='cuda' if torch.cuda.is_available() else 'cpu')

#generate(pth=)



